# -*- coding: utf-8 -*-
import numpy as np
import cv2
import glob
import os





global SAVE_PATH

SAVE_PATH = 'CorrectFiles'

if not os.path.exists(SAVE_PATH):
    os.mkdir(SAVE_PATH)

def calibrate(boardSize, squareSize, path):
    '''
        boardSize: 棋盘格的横纵方向上格子数
        squareSize: 棋盘格的大小
        path: 存放左右图片的根文件夹
    '''
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    Nx_cor = boardSize[0]
    Ny_cor = boardSize[1]

    objp = np.zeros((Nx_cor * Ny_cor, 3), np.float32)
    objp[:, :2] = np.mgrid[0:Nx_cor * squareSize:squareSize, 0:Ny_cor * squareSize:squareSize].T.reshape(-1, 2)# objp是三维的
    objpoints = []  # 3d points in real world space
    imgpoints = [[], []]  # 2d points in image plane.
    print(os.path.join(path, 'left'))
    left_images = glob.glob(r'{}/*.jpg'.format(os.path.join(path, 'left')))
    print(left_images)
    imageShape = None
    for fname_left in left_images:
        frame_left = cv2.imread(fname_left)
        fname_right = os.path.join(path, os.path.join('right', os.path.basename(fname_left)))
        if not os.path.exists(fname_right):
            print("{} not in right".format(os.path.basename(fname_right)))
            continue
        frame_right = cv2.imread(fname_right)
        # Our operations on the frame come here
        if not frame_left.shape == frame_right.shape:
            print('{} has different shape in left and right'.format(os.path.basename(fname_right)))
            continue
        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY)
        if not imageShape:
            imageShape = gray_left.shape
        elif gray_left.shape != imageShape:
            print('{} has different shape with before'.format(os.path.basename(fname_right)))
            continue
        print('Start Find Chess\n')
        ret_left, corners_left = cv2.findChessboardCorners(gray_left, (Nx_cor, Ny_cor), None)  # Find the corners
        ret_right, corners_right = cv2.findChessboardCorners(gray_right, (Nx_cor, Ny_cor), None)
        
        # If found, add object points, image points
        print('End Find Chess\n')
        if ret_left and ret_right:
            print('Start compute super pix\n')
            corners_left = cv2.cornerSubPix(gray_left, corners_left, (5, 5), (-1, -1), criteria)
            corners_right = cv2.cornerSubPix(gray_right, corners_right, (5, 5), (-1, -1), criteria)
            objpoints.append(objp)
            imgpoints[0].append(corners_left)
            imgpoints[1].append(corners_right)
            # show corners
            # cv2.drawChessboardCorners(frame_left, (Nx_cor, Ny_cor), corners_left, ret_left)
            # cv2.imshow('frame', frame_left)
            # k = cv2.waitKey(0)
            # if k == 27:
            #     cv2.destroyAllWindows()
        else:
            continue
        print('Finish {}\n'.format(fname_left))

    
    global mtx_left, dist_left, mtx_right, dist_right, R1, P1, R2, P2, Q, R, T, shape
    shape = imageShape
    ret_left, mtx_left, dist_left, rvecs_left, tvecs_left = cv2.calibrateCamera(objpoints, imgpoints[0], gray_left.shape[::-1], None, None)


    mean_error = 0
    for i in range(len(objpoints)):
        imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs_left[i], tvecs_left[i], mtx_left, dist_left)
        error = cv2.norm(imgpoints[0][i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        mean_error += error

    print("total error left: ", mean_error / len(objpoints))

    ret_right, mtx_right, dist_right, rvecs_right, tvecs_right = cv2.calibrateCamera(objpoints, imgpoints[1],
                                                                                gray_right.shape[::-1], None, None)

    mean_error = 0
    for i in range(len(objpoints)):
        imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs_right[i], tvecs_right[i], mtx_right, dist_right)
        error = cv2.norm(imgpoints[1][i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        mean_error += error

    print("total error right: ", mean_error / len(objpoints))
    print('mtx_left is {}'.format(mtx_left))
    print('dist_left {}'.format(dist_left))
    print('mtx_right is {}'.format(mtx_right))
    print('dist_right is {}'.format(dist_right))
    # # When everything done, release the capture

    # w, h, mtx_left = getNewMTX(mtx_left, dist_left)
        
    stereocalibration_criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 500, 1e-5)
    stereocalibration_flags = cv2.CALIB_FIX_INTRINSIC
    stereocalibration_retval, mtx_left, dist_left, mtx_right, dist_right, R, T, E, F = cv2.stereoCalibrate(objpoints, imgpoints[0], imgpoints[1], mtx_left, dist_left, mtx_right, dist_right, gray_left.shape[::-1], criteria=stereocalibration_criteria, flags=stereocalibration_flags)
    print('R is {}'.format(R))
    print('T is {}'.format(T))
    global P1, P2, Q
    if stereocalibration_retval:
        R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(mtx_left, dist_left, mtx_right, dist_right, gray_left.shape, R, T)
        np.savez('calibrate.npz', mtx_left=mtx_left, dist_left=dist_left,
                                  mtx_right=mtx_right, dist_right=dist_right,
                                  R=R, T=T, R1=R1, R2=R2, P1=P1, P2=P2, Q=Q,
                                  shape=shape)
    
    print('P1 is {}'.format(P1))
    print('P2 is {}'.format(P2))
    print('Q is {}'.format(Q))
    # cv2.triangulatePoints()

def getNewMTX(mtx, dist):
    h, w = shape
    newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
    x ,y, w, h = roi
    return w, h, newcameramtx


def correctImage(imageL, imageR, shape, mtx_left, dist_left, R1, mtx_right, dist_right, R2):
    
    imgL = cv2.imread(imageL)
    imgR = cv2.imread(imageR)
    h, w = shape
    newcameramtx_left, roi_left=cv2.getOptimalNewCameraMatrix(mtx_left,dist_left,(w,h),1,(w,h))
    # dst = cv2.undistort(imgL, mtx_left, dist_left, R1, newcameramtx_left)
    mapx,mapy = cv2.initUndistortRectifyMap(mtx_left,dist_left,R1, newcameramtx_left,(w,h),5)
    # mapx,mapy = cv2.initUndistortRectifyMap(mtx_left,dist_left,R1, newcameramtx_left,(w,h),5)
    dst_left = cv2.remap(imgL,mapx,mapy,cv2.INTER_LINEAR)
    
    # crop the image
    # x_l,y_l,w_l,h_l = roi_left
    # dst_left = dst_left[y_l:y_l+h_l, x_l:x_l+w_l]
    
    
    newcameramtx_right, roi_right=cv2.getOptimalNewCameraMatrix(mtx_right,dist_right,(w,h),1,(w,h))
    # dst_right = cv2.undistort(imgR, mtx_right, dist_right, R2, newcameramtx_right)
    mapx,mapy = cv2.initUndistortRectifyMap(mtx_right,dist_right,R2,newcameramtx_right,(w,h),5)
    # mapx,mapy = cv2.initUndistortRectifyMap(mtx_right,dist_right,R2,newcameramtx_right,(w,h),5)
    dst_right = cv2.remap(imgR,mapx,mapy,cv2.INTER_LINEAR)

    # crop the image
    # x,y,w,h = roi_right
    # dst_right = dst_right[y:y+h, x:x+w]
    
    cv2.imwrite(os.path.join(SAVE_PATH, 'left2_new1.jpg'), dst_left)
    cv2.imwrite(os.path.join(SAVE_PATH, 'right2_new1.jpg'), dst_right)

def detectPoint():
    fnameL = os.path.join(SAVE_PATH, 'left.jpg')
    fnameR = os.path.join(SAVE_PATH, 'right.jpg')
    if not os.path.exists(fnameL) or not os.path.exists(fnameR):
        print("The image is not exist")
        return

    imgL = cv2.imread(fnameL)
    imgR = cv2.imread(fnameR)
    points = []
    pass
    return points



def Clc3D(lpnt, rpnt, mtx_left, Q):
    '''
        lpnt:需要检测的左边点, 2*n numpy
        rpnt:需要检测的右边点, 2*n numpy
    '''
     #    point4D = cv2.triangulatePoints(P1, P2, lpnt, rpnt)
    d = lpnt[0] - rpnt[0]
    f = mtx_left[0, 0]
    w = Q[3, 2] * d + Q[3,3]
    
    x = -(lpnt[0] + Q[0,3])/ w
    y = -(lpnt[1] + Q[1,3]) / w
    z = -f / w
    return [x, y, z]
     # return point4D


def main(width, height, squareSize, path, target_path):
    '''
        width: 棋盘格的水平方向个数
        height: 棋盘格的竖直方向个数
        squareSize: 棋盘格在现实中的大小
        path: 照片的存放文件夹，文件夹内left为左边相机的照片，right为右边相机的照片
        target_path: 存放待检测图片的文件夹
    '''
    boardSize = (width, height)

    try:
        npzfile = np.load('calibrate.npz')
        mtx_left = npzfile['mtx_left']
        dist_left = npzfile['dist_left']
        mtx_right = npzfile['mtx_right']
        dist_right = npzfile['dist_right']
        R = npzfile['R']
        T = npzfile['T']
        R1 = npzfile['R1']
        # T1 = npzfile['T1']
        P1 = npzfile['P1']
        R2 = npzfile['R2']
        # T2 = npzfile['T2']
        P2 = npzfile['P2']
        T = npzfile['T']
        Q = npzfile['Q']
        shape = npzfile['shape']
    except IOError:
        calibrate(boardSize, squareSize, path)
    correctImage(r'20200103Pattern2/target_left.jpg', '20200103Pattern2/target_right.jpg', shape, mtx_left, dist_left, R1, mtx_right, dist_right, R2)
    # 将需要识别的图片进行矫正之后，确定目标的位置（即下面的lpnt和rpnt分别代表在左图的坐标和右图的坐标），再进行深度的求解
    lpnt = np.array([2062,1631 ]).reshape(-1, 1)
    rpnt = np.array([550,1628 ]).reshape(-1, 1)
    print(Clc3D(lpnt, rpnt, mtx_left, Q))
#    correctImage(os.path.join(target_path, 'left.jpg'), os.path.join(target_path), 'right.jpg')
#    points = detectPoint()
#    point3D = Clc3D(points[0], points[1])

if __name__ == '__main__':
    main(9, 6, 2.65, r'文件名', None)# 文件下包含left和right两个文件夹，文件夹内存放左边和右边相机拍摄的图片，注意图片的名称应一致

    
    